"""
1、当历史记录token 数量超过 max_token 则触发摘要总结
2、summary成员变量用于存储摘要
3、chat_histories 用于存储历史对话
4、get_num_tokens 用于计算对话的token 数量
5、save_context 用于存储新一轮对话到历史对话中
5、get_buffer_string 用于将历史对话转换成字符串
7、load_memory_variables 用于加载记忆变量信息
8、summary_text 用于将旧的摘要和传入的对话生成新的摘要
"""
import dotenv
from openai import OpenAI
from typing import Any


class ConversationSummaryBufferMemory:
    def __init__(self,summary = "", chat_histories:list=None,max_tokens:int = 300):
        self.m_summary = summary
        self.m_chat_histories = chat_histories
        self.m_max_tokens =max_tokens
        self.m_client = OpenAI(base_url="https://api.ephone.chat/v1")

    def summary_text(self,old_summary:str, new_line:str)-> str:
        """用于将旧的摘要和传入的对话生成新的摘要"""
        prompt = f"""
        你是一个强大的聊天机器人，请根据用户提供的谈话内容，总结摘要，并将其添加到先前提供的摘要中，返回一个新的摘要，除了新摘要其他任何数据都不要生成，如果用户的对话信息里有一些关键的信息，比方说姓名、爱好、性别、重要事件等等，这些全部都要包括在生成的摘要中，摘要尽可能要还原用户的对话记录。
        请不要将<example>标签里的数据当成实际的数据，这里的数据只是一个示例数据，告诉你该如何生成新摘要。
        <example>
        当前摘要：人类会问人工智能对人工智能的看法，人工智能认为人工智能是一股向善的力量。
        新的对话：
        Human：为什么你认为人工智能是一股向善的力量？
        AI：因为人工智能会帮助人类充分发挥潜力。
        新摘要：人类会问人工智能对人工智能的看法，人工智能认为人工智能是一股向善的力量，因为它将帮助人类充分发挥潜力。
        </example>
        =====================以下的数据是实际需要处理的数据=====================
        当前摘要：{old_summary}
        新的对话：
        {new_line}
        请帮用户将上面的信息生成新摘要。
        """
        completion = self.m_client.chat.completions.create(
            model = "gpt-3.5-turbo-16k",
            messages = [{"role":"user", "content":prompt}]

        )
        return completion.choices[0].message.content

    def get_buffer_string(self)->str:
        """获取所有聊天记录的字符串"""
        buffer : str =""
        for chat in self.m_chat_histories:
            buffer += f"Human:{chat.get('human')}\n AI:{chat.get('ai')}\n\n"
        return buffer


    def load_memory_variables(self)->str:
        """用于加载记忆变量信息"""
        buffer_string = self.get_buffer_string()
        return {
            "chat_histories":f"摘要:{self.m_summary}\n\n 历史信息:{buffer_string}"
        }
    @classmethod
    def get_num_tokens(self,query:str)->int:
        """用于计算对话的token 数量"""
        return int(len(query)*1.5)
        pass

    def save_context(self,human_query:str,ai_content:str):
        """用于存储新一轮对话到历史对话中"""
        self.m_chat_histories.append({"human":human_query, "ai":ai_content})
        buffer_string = self.get_buffer_string()
        tokens = self.get_num_tokens(buffer_string)

        if tokens>self.m_max_tokens:
            first_chat = self.m_chat_histories[0]
            print("==========新的摘要生成中。。。===========")
            self.m_summary = self.summary_text(self.m_summary,f"Human:{first_chat.get('human')}\n AI:{first_chat.get('ai')}")
            print(f"新摘要：{self.m_summary}")
            print("============新摘要生成完毕！！！=============")
            del self.m_chat_histories[0]

# 写一个简单的聊天机器人聊一天
dotenv.load_dotenv()
chat_client = OpenAI(base_url="https://api.ephone.chat/v1")

if __name__ == "__main__":
    app_mempry = ConversationSummaryBufferMemory("",[],300)
    while True:
        human_query = input("Human:")

        if human_query== "exit":
            print("Goodbye Human")
            exit(0)
        memory_variables = app_mempry.load_memory_variables()
        print(memory_variables)
        anwser_prompt = (
            "你是一个强大的聊天机器人，请根据对应的用户上下文和用户的问题来回答！\n\n"
            f"{memory_variables.get('chat_histories')}\n"
            f"用户的提问是：{human_query}"
        )
        completion = chat_client.chat.completions.create(
            model = "gpt-3.5-turbo-16k",
            messages = [{"role":"system","content":"你是OpenAi研发的机器人，请根据用户问题来回答问题！"},
                        {"role":"user","content":anwser_prompt}],
            stream = True
        )
        # print(completion)
        ai_content = ""
        print("AI",end="",flush=True)
        for chunk in completion:
            if len(chunk.choices) > 0:
                one_content = chunk.choices[0].delta.content
                if one_content is None:
                    break
                print(one_content,end="",flush=True)
                ai_content += one_content
        print("\n")
        app_mempry.save_context(human_query, ai_content)
